// Copyright (C) 2024 Kumo inc.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//

#pragma once

#include <llama.h>
#include <vector>
#include <string>

namespace kllm {

    enum ErrorType {
        ERROR_TYPE_INVALID_REQUEST,
        ERROR_TYPE_AUTHENTICATION,
        ERROR_TYPE_SERVER,
        ERROR_TYPE_NOT_FOUND,
        ERROR_TYPE_PERMISSION,
        ERROR_TYPE_UNAVAILABLE, // custom error
        ERROR_TYPE_NOT_SUPPORTED, // custom error
    };

    enum StopType {
        STOP_TYPE_FULL,
        STOP_TYPE_PARTIAL,
    };

    enum SlotState {
        SLOT_STATE_IDLE,
        SLOT_STATE_STARTED, // TODO: this state is only used for setting up the initial prompt processing; maybe merge it with launch_slot_with_task in the future
        SLOT_STATE_PROCESSING_PROMPT,
        SLOT_STATE_DONE_PROMPT,
        SLOT_STATE_GENERATING,
    };

    enum ServerState {
        SERVER_STATE_LOADING_MODEL,  // Server is starting up, model not fully loaded yet
        SERVER_STATE_READY,          // Server is ready and model is loaded
    };

    enum ServerTaskType {
        SERVER_TASK_TYPE_INFERENCE,
        SERVER_TASK_TYPE_CANCEL,
        SERVER_TASK_TYPE_NEXT_RESPONSE,
        SERVER_TASK_TYPE_METRICS,
        SERVER_TASK_TYPE_SLOT_SAVE,
        SERVER_TASK_TYPE_SLOT_RESTORE,
        SERVER_TASK_TYPE_SLOT_ERASE,
        SERVER_TASK_TYPE_SET_LORA,
    };

    enum ServerTaskInfType {
        SERVER_TASK_INF_TYPE_COMPLETION,
        SERVER_TASK_INF_TYPE_EMBEDDING,
        SERVER_TASK_INF_TYPE_RERANK,
        SERVER_TASK_INF_TYPE_INFILL,
    };

    struct completion_token_output {
        llama_token tok;
        std::string text_to_send;

        struct token_prob {
            llama_token tok;
            float prob;
        };

        std::vector<token_prob> probs;
    };
}  // namespace kllm
